//! Interpreter for stencil read operations using the AST.

use thiserror::Error;

use crate::mir::{EnvExpr, Expr};
use crate::tysys::{ArgRef, BinOp, FieldRef, LiteralValue, MonOp, ValueType};

#[derive(Debug, Error)]
pub enum Error {
    /// Explicit failure from the expr being evaluated.  This should only happen
    /// when checking a buffer when wrapping it in a type.
    #[error("expr failed (code {0})")]
    ExprFail(u32),

    /// Tried to read offset out of bounds.
    #[error("out of bounds read (at {0})")]
    OutOfBounds(usize),

    /// Tried to do a math out of an integer range.
    #[error("integer operation out of range")]
    OutOfRange,

    /// Generic type interpretation failure.
    #[error("type mismatch (expected {0:?}")]
    TypeMismatch(ValueType),

    /// Impossible conversion. (dest, source)
    #[error("invalid cast (from {1:?} to {0:?})")]
    InvalidCast(ValueType, ValueType),

    /// Impossible conversion from buffer. (dest, source len)
    #[error("invalid buffer cast (from buf of size {1} to {0:?})")]
    InvalidBufCast(ValueType, usize),

    /// Type mismatch on a mono-op.
    #[error("type mismatch on {0:?} (from {1:?})")]
    TypeMismatchMonop(MonOp, ValueType),

    /// Type mismatch on a binary operation.
    #[error("type mismatch on {0:?} (from {1:?} and {2:?})")]
    TypeMismatchBinop(BinOp, ValueType, ValueType),

    /// Tried to write to an immutable backing buffer.
    #[error("write in immutable context")]
    Immutable,

    /// Unknown field.
    #[error("unknown field {0:?}")]
    UnknownField(FieldRef),

    /// Unknown argument.
    #[error("unknown argument {0:?}")]
    UnknownArg(ArgRef),

    /// Not yet implemented.
    #[error("not yet implemented")]
    Unimplemented,
}

/// Values we can extract from the source.
#[derive(Clone, Debug)]
pub enum Value<'b> {
    Void,
    U8(u8),
    U64(u64),
    I64(i64),

    /// Use a usize here because we're actually interpreting here.
    USize(usize),

    /// Same as in `USize`.
    ISize(isize),
    Bool(bool),
    Slice(&'b [u8]),
}

impl<'b> Value<'b> {
    pub fn from_literal(lv: &LiteralValue) -> Self {
        match lv {
            LiteralValue::Void => Self::Void,
            LiteralValue::U8(v) => Self::U8(*v),
            LiteralValue::U64(v) => Self::U64(*v),
            LiteralValue::I64(v) => Self::I64(*v),
            LiteralValue::USize(v) => Self::USize(*v as usize),
            LiteralValue::ISize(v) => Self::ISize(*v as isize),
            LiteralValue::Bool(b) => Self::Bool(*b),
        }
    }

    pub fn ty(&self) -> ValueType {
        match self {
            Self::Void => ValueType::Void,
            Self::U8(_) => ValueType::U8,
            Self::U64(_) => ValueType::U64,
            Self::I64(_) => ValueType::I64,
            Self::USize(_) => ValueType::USize,
            Self::ISize(_) => ValueType::ISize,
            Self::Bool(_) => ValueType::Bool,
            Self::Slice(_) => ValueType::ByteSlice,
        }
    }
}

pub struct RunContext<'c> {
    fields: Vec<Value<'c>>,
    args: Vec<Value<'c>>,
}

impl<'c> RunContext<'c> {
    pub fn empty() -> Self {
        Self {
            fields: Vec::new(),
            args: Vec::new(),
        }
    }

    /// Adds a field with some value, returning an expr that evaluates to it.
    pub fn add_field(&mut self, v: Value<'c>) -> Expr {
        let i = self.fields.len() as u16;
        self.fields.push(v);
        Expr::new_field(FieldRef::from(i))
    }

    /// Adds an arg with some value, returning an expr that evaluates to it.
    pub fn add_arg(&mut self, v: Value<'c>) -> Expr {
        let i = self.args.len() as u16;
        self.args.push(v);
        Expr::new_arg(ArgRef::from(i))
    }

    pub fn resolve_env(&self, expr: &EnvExpr) -> Result<Value<'c>, Error> {
        match expr {
            EnvExpr::Field(i) => self
                .fields
                .get(i.inner() as usize)
                .cloned()
                .ok_or(Error::UnknownField(*i)),
            EnvExpr::Arg(i) => self
                .args
                .get(i.inner() as usize)
                .cloned()
                .ok_or(Error::UnknownArg(*i)),
            EnvExpr::Local(_i) => Err(Error::Unimplemented),
        }
    }
}

pub struct Interpreter<'b> {
    run_context: RunContext<'b>,
}

impl<'b> Interpreter<'b> {
    pub fn new(ctx: RunContext<'b>) -> Self {
        Self { run_context: ctx }
    }

    pub fn get_field(&self, field: FieldRef) -> Result<Value<'b>, Error> {
        self.run_context
            .fields
            .get(field.inner() as usize)
            .cloned()
            .ok_or(Error::UnknownField(field))
    }

    pub fn get_arg(&self, arg: ArgRef) -> Result<Value<'b>, Error> {
        self.run_context
            .args
            .get(arg.inner() as usize)
            .cloned()
            .ok_or(Error::UnknownArg(arg))
    }

    pub fn compute(&self, expr: &Expr) -> Result<Value<'b>, Error> {
        let v = match expr {
            Expr::Const(_, e) => Value::from_literal(e),

            // TODO implement fields and args in the interpreter context
            Expr::Env(_, e) => self.run_context.resolve_env(e)?,

            Expr::Monop(_, mo, e) => {
                use MonOp::*;
                use Value::*;

                let v = self.compute(e)?;
                match (mo, v) {
                    (Neg, I64(v)) => {
                        if v == i64::MIN {
                            return Err(Error::OutOfRange);
                        }
                        I64(-v)
                    }
                    (Neg, U8(_) | U64(_)) => return Err(Error::OutOfRange),
                    (Neg, Bool(v)) => Bool(!v),
                    (BNot, U8(v)) => U8(!v),
                    (BNot, I64(v)) => I64(!v),
                    (BNot, U64(v)) => U64(!v),
                    (BNot, Bool(v)) => Bool(!v),
                    (mo, v) => return Err(Error::TypeMismatchMonop(*mo, v.ty())),
                }
            }

            Expr::Binop(_, bo, lhs, rhs) => {
                use BinOp::*;
                use Value::*;

                // TODO make all this crap macro-generatable
                let lhsv = self.compute(lhs)?;
                let rhsv = self.compute(rhs)?;
                match (bo, lhsv, rhsv) {
                    // Boolean comparisons.
                    (Eq, Bool(lv), Bool(rv)) => Bool(lv == rv),
                    (Ne, Bool(lv), Bool(rv)) => Bool(lv != rv),

                    // Buffer comparisons.
                    (Eq, Slice(lv), Slice(rv)) => Bool(lv == rv),
                    (Ne, Slice(lv), Slice(rv)) => Bool(lv != rv),

                    // Integer math, bitwise ops.
                    (Eq, U8(lv), U8(rv)) => Bool(lv == rv),
                    (Ne, U8(lv), U8(rv)) => Bool(lv != rv),
                    (Add, U8(lv), U8(rv)) => U8(lv + rv),
                    (Sub, U8(lv), U8(rv)) => U8(lv - rv),
                    (Mul, U8(lv), U8(rv)) => U8(lv * rv),
                    (Shr, U8(lv), U8(rv)) => U8(lv >> rv),
                    (Shl, U8(lv), U8(rv)) => U8(lv << rv),
                    (BAnd, U8(lv), U8(rv)) => U8(lv & rv),
                    (BOr, U8(lv), U8(rv)) => U8(lv | rv),

                    (Eq, U64(lv), U64(rv)) => Bool(lv == rv),
                    (Ne, U64(lv), U64(rv)) => Bool(lv != rv),
                    (Add, U64(lv), U64(rv)) => U64(lv + rv),
                    (Sub, U64(lv), U64(rv)) => U64(lv - rv),
                    (Mul, U64(lv), U64(rv)) => U64(lv * rv),
                    (Shr, U64(lv), U64(rv)) => U64(lv >> rv),
                    (Shl, U64(lv), U64(rv)) => U64(lv << rv),
                    (BAnd, U64(lv), U64(rv)) => U64(lv & rv),
                    (BOr, U64(lv), U64(rv)) => U64(lv | rv),

                    (Eq, I64(lv), I64(rv)) => Bool(lv == rv),
                    (Ne, I64(lv), I64(rv)) => Bool(lv != rv),
                    (Add, I64(lv), I64(rv)) => I64(lv + rv),
                    (Sub, I64(lv), I64(rv)) => I64(lv - rv),
                    (Mul, I64(lv), I64(rv)) => I64(lv * rv),
                    (Shr, I64(lv), I64(rv)) => I64(lv >> rv),
                    (Shl, I64(lv), I64(rv)) => I64(lv << rv),
                    (BAnd, I64(lv), I64(rv)) => I64(lv & rv),
                    (BOr, I64(lv), I64(rv)) => I64(lv | rv),

                    // TODO more
                    (Eq, USize(lv), USize(rv)) => Bool(lv == rv),
                    (Eq, ISize(lv), ISize(rv)) => Bool(lv == rv),

                    // All other combinations are invalid.
                    (bo, lv, rv) => return Err(Error::TypeMismatchBinop(*bo, lv.ty(), rv.ty())),
                }
            }

            Expr::Cond(_, c, a, b) => {
                let cv = self.compute(c)?;
                let Value::Bool(c) = cv else {
                    return Err(Error::TypeMismatch(cv.ty()));
                };

                if c {
                    self.compute(a)?
                } else {
                    self.compute(b)?
                }
            }

            Expr::Cast(_, dty, e) => {
                use Value::*;

                // TODO make all this crap macro-generatable
                let v = self.compute(e)?;

                match (dty, v) {
                    // Conversions between integers.
                    (ValueType::U8, U64(v)) => U8(v as u8),
                    (ValueType::U8, I64(v)) => U8(v as u8),
                    (ValueType::I64, U8(v)) => I64(v as i64),
                    (ValueType::I64, U64(v)) => I64(v as i64),
                    (ValueType::I64, Bool(b)) => I64(if b { 1 } else { 0 }),
                    (ValueType::U64, U8(v)) => U64(v as u64),
                    (ValueType::U64, I64(v)) => U64(v as u64),
                    (ValueType::U64, Bool(b)) => U64(if b { 1 } else { 0 }),
                    (ValueType::Bool, U8(v)) => Bool(v != 0),
                    (ValueType::Bool, U64(v)) => Bool(v != 0),
                    (ValueType::Bool, I64(v)) => Bool(v != 0),

                    // Conversions to integers from buffers.
                    (ValueType::U8, Slice(buf)) => {
                        if buf.len() != 0 {
                            return Err(Error::InvalidBufCast(dty.clone(), buf.len()));
                        }

                        U8(buf[0])
                    }

                    (ValueType::U64, Slice(buf)) => {
                        if buf.len() != 8 {
                            return Err(Error::InvalidBufCast(dty.clone(), buf.len()));
                        }

                        let mut db = [0; 8];
                        db.copy_from_slice(buf);
                        U64(u64::from_be_bytes(db))
                    }

                    (ValueType::I64, Slice(buf)) => {
                        if buf.len() != 8 {
                            return Err(Error::InvalidBufCast(dty.clone(), buf.len()));
                        }

                        let mut db = [0; 8];
                        db.copy_from_slice(buf);
                        I64(i64::from_be_bytes(db))
                    }

                    (ValueType::USize, U8(v)) => USize(v as usize),

                    (dty, v) => return Err(Error::InvalidCast(v.ty(), dty.clone())),
                }
            }

            Expr::BufferLen(_, bufe) => {
                let bufv = self.compute(bufe)?;

                let Value::Slice(buf) = bufv else {
                    return Err(Error::TypeMismatch(bufv.ty()));
                };

                Value::U64(buf.len() as u64)
            }

            Expr::BufferRead(_, bufe, sizee, offe) => {
                let bufv = self.compute(bufe)?;
                let sizev = self.compute(sizee)?;
                let offv = self.compute(offe)?;

                let Value::Slice(buf) = bufv else {
                    return Err(Error::TypeMismatch(bufv.ty()));
                };

                let Value::U64(size) = sizev else {
                    return Err(Error::TypeMismatch(sizev.ty()));
                };

                let Value::U64(off) = offv else {
                    return Err(Error::TypeMismatch(offv.ty()));
                };

                if off as usize >= buf.len() {
                    return Err(Error::OutOfBounds(off as usize));
                }

                let end = off + size;
                if end as usize > buf.len() {
                    return Err(Error::OutOfBounds(end as usize));
                }

                Value::Slice(&buf[off as usize..end as usize])
            }

            Expr::BufferReadInt(_, bufe, offe, _ienc) => {
                let bufv = self.compute(bufe)?;
                let offv = self.compute(offe)?;

                let Value::Slice(buf) = bufv else {
                    return Err(Error::TypeMismatch(bufv.ty()));
                };

                let Value::USize(off) = offv else {
                    return Err(Error::TypeMismatch(offv.ty()));
                };

                if off as usize >= buf.len() {
                    return Err(Error::OutOfBounds(off as usize));
                }

                // TODO implement bigger numbers
                Value::U8(buf[off as usize])
            }
        };

        Ok(v)
    }
}

#[cfg(test)]
mod tests {
    use crate::interpreter::{Interpreter, Value};
    use crate::{mir::*, test_util};

    use super::RunContext;

    #[test]
    fn test_interp_basic_1() {
        let buf = [1, 2, 3, 4, 5];
        let bufv = Value::Slice(&buf);
        let mut rc = RunContext::empty();
        let bufe = rc.add_field(bufv);
        let e = Expr::new_read_u8_at(bufe, 2);

        let interp = Interpreter::new(rc);

        let res = interp.compute(&e);
        match res {
            Ok(Value::U8(3)) => {}
            _ => panic!("bad result: {res:?}"),
        }
    }

    #[test]
    fn test_interp_cond() {
        let mut buf = [0, 5, 7];

        {
            let bufv = Value::Slice(&buf);
            let mut rc = RunContext::empty();
            let bufe = rc.add_field(bufv);
            let e = test_util::get_012_cond_expr(bufe);

            let interp = Interpreter::new(rc);
            let res = interp.compute(&e);
            match res {
                Ok(Value::U8(7)) => {}
                _ => panic!("bad result: {res:?}"),
            }
        }

        buf[0] = 2;

        {
            let bufv = Value::Slice(&buf);
            let mut rc = RunContext::empty();
            let bufe = rc.add_field(bufv);
            let e = test_util::get_012_cond_expr(bufe);

            let interp = Interpreter::new(rc);
            let res = interp.compute(&e);
            match res {
                Ok(Value::U8(5)) => {}
                _ => panic!("bad result: {res:?}"),
            }
        }
    }
}
